import json
import secrets
from datetime import datetime
import requests
from cozepy import load_oauth_app_from_config, JWTOAuthApp
from flask import Flask, redirect, session, request, jsonify

app = Flask(
    __name__,
    static_folder="assets",  # use shared/assets as static directory
    static_url_path="/assets",
)  # URL path for static files
app.secret_key = secrets.token_hex(16)  # for Flask session encryption

COZE_OAUTH_CONFIG_PATH = "coze_oauth_config.json"
REDIRECT_URI = "http://127.0.0.1:8093/callback"


def load_app_config(config_path) -> dict:
    with open(config_path, "r") as file:
        config = file.read()
    return json.loads(config)


def load_coze_oauth_app(config_path) -> JWTOAuthApp:
    try:
        with open(config_path, "r") as file:
            config = json.loads(file.read())
        return load_oauth_app_from_config(config)  # type: ignore
    except FileNotFoundError:
        raise Exception(
            f"Configuration file not found: {config_path}. Please make sure you have created the OAuth configuration file."
        )
    except Exception as e:
        raise Exception(f"Failed to load OAuth configuration: {str(e)}")


def timestamp_to_datetime(timestamp: int) -> str:
    return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S")


app_config = load_app_config(COZE_OAUTH_CONFIG_PATH)
coze_oauth_app = load_coze_oauth_app(COZE_OAUTH_CONFIG_PATH)


def read_html_template(file_path: str) -> str:
    with open(file_path, "r", encoding="utf-8") as file:
        return file.read()


def render_template(template: str, kwargs: dict) -> str:
    if not kwargs:
        kwargs = {}
    kwargs["coze_www_base"] = app_config["coze_www_base"]
    template = read_html_template(template)
    for key, value in kwargs.items():
        template = template.replace(f"{{{{{key}}}}}", str(value))
    return template


@app.errorhandler(Exception)
def handle_error(error):
    error_message = str(error)
    print(f"Error occurred: {error_message}")
    return render_template("websites/error.html", {"error": error_message})


@app.route("/")
def index():
    if not coze_oauth_app:
        return render_template(
            "websites/error.html",
            {
                "error": "OAuth application is not properly configured. Please check your configuration file."
            },
        )
    return render_template("websites/index.html", app_config)


@app.route("/login")
def login():
    return redirect("/callback")


@app.route("/callback")
def callback():
    if not coze_oauth_app:
        return render_template(
            "websites/error.html",
            {
                "error": "OAuth application is not properly configured. Please check your configuration file."
            },
        )

    try:
        oauth_token = coze_oauth_app.get_access_token()
        # 将 OAuth token 保存到 session 中，以便后续使用
        session[f'oauth_token_{app_config["client_id"]}'] = {
            "token_type": oauth_token.token_type,
            "access_token": oauth_token.access_token,
            "expires_in": oauth_token.expires_in,
        }

        expires_str = timestamp_to_datetime(oauth_token.expires_in)

        # 如果是 AJAX 请求，返回 JSON 格式
        if request.headers.get("X-Requested-With") == "XMLHttpRequest":
            return {
                "token_type": oauth_token.token_type,
                "access_token": oauth_token.access_token,
                "refresh_token": "",
                "expires_in": f"{oauth_token.expires_in} ({expires_str})",
            }

        # 否则返回 HTML 页面
        return render_template(
            "websites/callback.html",
            {
                "token_type": oauth_token.token_type,
                "access_token": oauth_token.access_token,
                "refresh_token": "",
                "expires_in": f"{oauth_token.expires_in} ({expires_str})",
            },
        )
    except Exception as e:
        error_message = f"Failed to get access token: {str(e)}"
        if request.headers.get("X-Requested-With") == "XMLHttpRequest":
            return {"error": error_message}, 500
        return render_template("websites/error.html", {"error": error_message})


# 新增聊天页面路由
@app.route("/chat")
def chat_page():
    return render_template("websites/chat.html", {})


# 新增API端点：获取令牌
@app.route("/api/token", methods=["POST"])
def get_token():
    try:
        # 从请求中获取会话名称
        data = request.json or {}
        session_name = data.get("sessionName", "default")
        
        # 为每个会话获取独立的令牌
        oauth_token = coze_oauth_app.get_access_token()
        
        # 将令牌保存到会话中
        session_key = f'oauth_token_{session_name}'
        session[session_key] = {
            "token_type": oauth_token.token_type,
            "access_token": oauth_token.access_token,
            "expires_in": oauth_token.expires_in,
        }
        
        # 返回令牌信息
        return jsonify({
            "token_type": oauth_token.token_type,
            "access_token": oauth_token.access_token,
            "expires_in": oauth_token.expires_in
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500


# 新增API端点：创建会话
@app.route("/api/conversation/create", methods=["POST"])
def create_conversation():
    try:
        # 获取请求数据
        data = request.json or {}
        session_id = data.get("session_id", "default")
        
        # 从会话中获取OAuth令牌
        session_key = f'oauth_token_{session_id}'
        oauth_data = session.get(session_key)
        
        # 如果会话中没有令牌，尝试获取新令牌
        if not oauth_data:
            try:
                oauth_token = coze_oauth_app.get_access_token()
                oauth_data = {
                    "token_type": oauth_token.token_type,
                    "access_token": oauth_token.access_token,
                    "expires_in": oauth_token.expires_in,
                }
                session[session_key] = oauth_data
            except Exception as e:
                print(f"获取OAuth令牌失败: {str(e)}")
                return jsonify({"error": f"获取OAuth令牌失败: {str(e)}"}), 500
        
        # 使用OAuth令牌
        access_token = oauth_data.get("access_token")
        if not access_token:
            return jsonify({"error": "无效的访问令牌"}), 401
        
        # 构建请求头
        headers = {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json"
        }
        
        # 创建会话
        conversation_api_url = "https://api.coze.cn/v1/conversation/create"
        
        print(f"创建会话请求: {conversation_api_url}")
        print(f"请求头: {headers}")
        
        # 发送创建会话请求 - 不需要请求体
        conversation_response = requests.post(
            conversation_api_url,
            headers=headers,
            timeout=60
        )
        
        # 检查会话创建响应
        if conversation_response.status_code != 200:
            print(f"创建会话失败: {conversation_response.status_code}")
            print(f"错误响应: {conversation_response.text}")
            return jsonify({
                "error": f"创建会话失败: {conversation_response.status_code}",
                "details": conversation_response.text
            }), conversation_response.status_code
        
        # 解析会话创建响应
        try:
            conversation_data = conversation_response.json()
            print(f"会话创建响应: {conversation_data}")
            
            # 检查响应码
            if conversation_data.get("code") != 0:
                error_msg = conversation_data.get("msg", "未知错误")
                return jsonify({"error": f"创建会话API错误: {error_msg}"}), 500
            
            # 获取会话ID
            conversation_id = conversation_data.get("data", {}).get("id")
            if not conversation_id:
                return jsonify({"error": "创建会话成功但未返回会话ID"}), 500
                
            print(f"成功创建会话，ID: {conversation_id}")
            
            # 返回会话ID
            return jsonify({
                "conversation_id": conversation_id,
                "session_id": session_id
            })
        except Exception as e:
            print(f"解析会话创建响应失败: {str(e)}")
            return jsonify({"error": f"解析会话创建响应失败: {str(e)}"}), 500
            
    except Exception as e:
        print(f"创建会话失败: {str(e)}")
        return jsonify({"error": f"创建会话失败: {str(e)}"}), 500

# 新增API端点：聊天
@app.route("/api/chat", methods=["POST"])
def chat():
    try:
        # 获取请求数据
        data = request.json
        if not data:
            return jsonify({"error": "无效的请求数据"}), 400
        
        # 获取会话ID和消息
        session_id = data.get("session_id", "default")
        messages = data.get("messages", [])
        conversation_id = data.get("conversation_id")
        workflow_id = data.get("workflow_id", "7486679564925829161")
        
        if not messages:
            return jsonify({"error": "缺少必要参数"}), 400
            
        # 从会话中获取OAuth令牌
        session_key = f'oauth_token_{session_id}'
        oauth_data = session.get(session_key)
        
        # 如果会话中没有令牌，尝试获取新令牌
        if not oauth_data:
            try:
                oauth_token = coze_oauth_app.get_access_token()
                oauth_data = {
                    "token_type": oauth_token.token_type,
                    "access_token": oauth_token.access_token,
                    "expires_in": oauth_token.expires_in,
                }
                session[session_key] = oauth_data
            except Exception as e:
                print(f"获取OAuth令牌失败: {str(e)}")
                return jsonify({"error": f"获取OAuth令牌失败: {str(e)}"}), 500
        
        # 使用OAuth令牌
        access_token = oauth_data.get("access_token")
        if not access_token:
            return jsonify({"error": "无效的访问令牌"}), 401
        
        # 构建请求头
        headers = {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json"
        }
        
        # 步骤1: 如果没有会话ID，先创建会话
        if not conversation_id:
            conversation_api_url = "https://api.coze.cn/v1/conversation/create"
            
            # 发送创建会话请求 - 简化版本，不需要请求体
            conversation_response = requests.post(
                conversation_api_url,
                headers=headers,
                timeout=60
            )
            
            # 检查会话创建响应
            if conversation_response.status_code != 200:
                print(f"创建会话失败: {conversation_response.status_code}")
                print(f"错误响应: {conversation_response.text}")
                return jsonify({
                    "error": f"创建会话失败: {conversation_response.status_code}",
                    "details": conversation_response.text
                }), conversation_response.status_code
            
            # 解析会话创建响应
            try:
                conversation_data = conversation_response.json()
                print(f"会话创建响应: {conversation_data}")
                
                # 检查响应码
                if conversation_data.get("code") != 0:
                    error_msg = conversation_data.get("msg", "未知错误")
                    return jsonify({"error": f"创建会话API错误: {error_msg}"}), 500
                
                # 获取会话ID
                conversation_id = conversation_data.get("data", {}).get("id")
                if not conversation_id:
                    return jsonify({"error": "创建会话成功但未返回会话ID"}), 500
                    
                print(f"成功创建会话，ID: {conversation_id}")
            except Exception as e:
                print(f"解析会话创建响应失败: {str(e)}")
                return jsonify({"error": f"解析会话创建响应失败: {str(e)}"}), 500
        
        # 格式化消息历史
        additional_messages = []
        for msg in messages:
            if isinstance(msg, dict) and "role" in msg and "content" in msg:
                additional_messages.append({
                    "role": msg["role"],
                    "content": msg["content"],
                    "content_type": "text"
                })
        
        # 步骤2: 使用workflows/chat接口进行对话
        chat_api_url = "https://api.coze.cn/v1/workflows/chat"
        
        # 构建chat请求体 - 添加 app_id
        chat_request = {
            "app_id": "7486666914035646505",  # 添加 app_id
            "additional_messages": additional_messages,
            "conversation_id": conversation_id,
            "workflow_id": workflow_id
        }
        
        print(f"发送请求到Coze API: {chat_api_url}")
        print(f"请求头: {headers}")
        print(f"请求体: {chat_request}")
        
        # 发送chat请求
        chat_response = requests.post(
            chat_api_url,
            headers=headers,
            json=chat_request,
            timeout=120  # 增加超时时间到120秒
        )
        
        # 检查响应状态码
        if chat_response.status_code != 200:
            print(f"Coze API返回错误状态码: {chat_response.status_code}")
            print(f"错误响应: {chat_response.text}")
            return jsonify({
                "error": f"Coze API返回错误: {chat_response.status_code}",
                "details": chat_response.text
            }), chat_response.status_code
        
        # 添加日志记录原始响应内容
        print(f"Coze API原始响应内容: {chat_response.text[:500]}...")  # 只打印前500个字符避免日志过长
        print(f"Coze API响应头: {chat_response.headers}")
        
        # 处理响应
        try:
            # 尝试解析JSON响应
            try:
                response_data = chat_response.json()
                print(f"Coze API返回JSON响应: {response_data}")
                
                # 检查响应码
                if response_data.get("code") != 0:
                    error_msg = response_data.get("msg", "未知错误")
                    return jsonify({"error": f"Coze API错误: {error_msg}"}), 500
                
                # 提取回复内容
                content = response_data.get("data", {}).get("content", "")
                if content:
                    return jsonify({
                        "content": content,
                        "role": "assistant",
                        "session_id": session_id,
                        "conversation_id": conversation_id
                    })
            except json.JSONDecodeError as e:
                # 如果不是JSON响应，可能是SSE格式
                print(f"响应不是JSON格式，尝试解析SSE格式。错误: {str(e)}")
            
            # 解析SSE响应
            sse_content = chat_response.text
            print(f"SSE响应长度: {len(sse_content)}")
            
            # 直接从SSE响应中提取完整的助手回复
            assistant_content = ""
            
            # 分割SSE事件
            events = [event for event in sse_content.split("\n\n") if event.strip()]
            print(f"SSE事件数量: {len(events)}")
            
            # 记录每个事件的类型
            event_types = []
            
            for i, event in enumerate(events):
                # 解析事件类型和数据
                event_lines = event.strip().split("\n")
                event_type = None
                event_data = None
                
                for line in event_lines:
                    if line.startswith("event:"):
                        event_type = line[6:].strip()
                        event_types.append(event_type)
                    elif line.startswith("data:"):
                        event_data = line[5:].strip()
                
                print(f"事件 {i}: 类型={event_type}, 数据长度={len(event_data) if event_data else 0}")
                
                # 如果是消息完成事件，直接提取内容
                if event_type == "conversation.message.completed" and event_data:
                    try:
                        # 解码JSON数据
                        data_json = json.loads(event_data)
                        print(f"消息完成事件数据: {data_json}")
                        if data_json.get("role") == "assistant":
                            # 获取内容并正确解码
                            content = data_json.get("content", "")
                            # 处理编码问题
                            try:
                                # 尝试不同的解码方式
                                if isinstance(content, bytes):
                                    content = content.decode('utf-8')
                                elif isinstance(content, str):
                                    # 如果内容是被错误编码的UTF-8字符串
                                    content = content.encode('latin1').decode('utf-8')
                                
                                # 移除可能存在的HTML标签
                                content = content.replace('<user-form></user-form>', '')
                                
                                assistant_content = content.strip()
                                print(f"从消息完成事件中获取回复（解码后）: {assistant_content}")
                            except Exception as decode_error:
                                print(f"解码内容失败: {decode_error}")
                                # 如果解码失败，尝试直接使用原始内容
                                assistant_content = content
                            break
                    except Exception as e:
                        print(f"解析消息完成事件失败: {str(e)}")
            
            print(f"所有事件类型: {event_types}")
            
            # 如果找到了助手回复，直接返回
            if assistant_content:
                print(f"返回助手回复: {assistant_content[:100]}...")  # 只打印前100个字符
                return jsonify({
                    "content": assistant_content,
                    "role": "assistant",
                    "session_id": session_id,
                    "conversation_id": conversation_id
                })
            
            # 如果没有找到消息完成事件，尝试提取聊天ID
            chat_id = None
            
            # 分割SSE事件
            for event in events:
                if not event.strip():
                    continue
                
                # 解析事件类型和数据
                event_lines = event.strip().split("\n")
                event_type = None
                event_data = None
                
                for line in event_lines:
                    if line.startswith("event:"):
                        event_type = line[6:].strip()
                    elif line.startswith("data:"):
                        event_data = line[5:].strip()
                
                # 如果是聊天创建事件，提取聊天ID
                if event_type == "conversation.chat.created" and event_data:
                    try:
                        data_json = json.loads(event_data)
                        chat_id = data_json.get("id")
                        print(f"找到聊天ID: {chat_id}")
                        break  # 找到ID后立即退出循环
                    except Exception as e:
                        print(f"解析聊天创建事件失败: {str(e)}")
            
            # 如果有chat_id，获取聊天内容
            if chat_id:
                print(f"开始轮询获取聊天ID {chat_id} 的结果")
                # 轮询获取聊天结果
                max_retries = 10  # 减少重试次数，避免超时
                retry_count = 0
                retry_interval = 2  # 轮询间隔时间（秒）
                
                while retry_count < max_retries:
                    # 等待指定时间
                    import time
                    time.sleep(retry_interval)
                    
                    # 获取聊天详情
                    chat_detail_url = f"https://api.coze.cn/v1/chat/{chat_id}"
                    print(f"轮询 {retry_count+1}/{max_retries}: 获取聊天详情 {chat_detail_url}")
                    
                    detail_response = requests.get(
                        chat_detail_url,
                        headers=headers,
                        timeout=30
                    )
                    
                    print(f"聊天详情响应状态码: {detail_response.status_code}")
                    
                    if detail_response.status_code == 200:
                        try:
                            detail_data = detail_response.json()
                            print(f"聊天详情响应: {detail_data}")
                            
                            # 检查响应码
                            if detail_data.get("code") != 0:
                                error_msg = detail_data.get("msg", "未知错误")
                                print(f"获取聊天详情失败: {error_msg}")
                                # 继续尝试，而不是立即返回错误
                                retry_count += 1
                                continue
                            
                            # 获取聊天状态
                            chat_info = detail_data.get("data", {})
                            status = chat_info.get("status")
                            print(f"聊天状态: {status}")
                            
                            # 如果聊天已完成，获取消息
                            if status == "completed":
                                # 获取消息
                                messages_url = f"https://api.coze.cn/v1/chat/{chat_id}/messages"
                                
                                print(f"获取聊天消息: {messages_url}")
                                messages_response = requests.get(
                                    messages_url,
                                    headers=headers,
                                    timeout=30
                                )
                                
                                print(f"聊天消息响应状态码: {messages_response.status_code}")
                                
                                if messages_response.status_code == 200:
                                    try:
                                        messages_data = messages_response.json()
                                        print(f"聊天消息响应: {messages_data}")
                                        
                                        # 检查响应码
                                        if messages_data.get("code") != 0:
                                            error_msg = messages_data.get("msg", "未知错误")
                                            print(f"获取聊天消息失败: {error_msg}")
                                            retry_count += 1
                                            continue
                                        
                                        # 获取消息列表
                                        message_list = messages_data.get("data", {}).get("messages", [])
                                        print(f"消息列表长度: {len(message_list)}")
                                        
                                        # 获取最后一条助手消息
                                        assistant_message = None
                                        for msg in reversed(message_list):
                                            if msg.get("role") == "assistant":
                                                assistant_message = msg
                                                break
                                        
                                        if assistant_message:
                                            content = assistant_message.get("content", "")
                                            print(f"助手回复: {content[:100]}...")  # 只打印前100个字符
                                            return jsonify({
                                                "content": content,
                                                "role": "assistant",
                                                "session_id": session_id,
                                                "conversation_id": conversation_id
                                            })
                                        else:
                                            print("未找到助手消息")
                                    except Exception as e:
                                        print(f"解析聊天消息失败: {str(e)}")
                            
                            # 如果聊天失败，返回错误
                            elif status == "failed":
                                error_info = chat_info.get("last_error", {})
                                error_code = error_info.get("code", 0)
                                error_msg = error_info.get("msg", "聊天处理失败")
                                
                                print(f"聊天处理失败: {error_msg}")
                                return jsonify({
                                    "error": f"聊天处理失败: {error_msg}",
                                    "code": error_code
                                }), 500
                            
                            # 如果聊天仍在进行中，继续等待
                            else:
                                print(f"聊天仍在进行中，继续等待...")
                                retry_count += 1
                                continue
                                
                        except Exception as e:
                            print(f"解析聊天详情失败: {str(e)}")
                            retry_count += 1
                            continue
                    else:
                        print(f"获取聊天详情失败: {detail_response.status_code}")
                        print(f"错误响应: {detail_response.text}")
                        retry_count += 1
                        continue
                
                # 如果达到最大重试次数仍未完成，返回部分结果
                print("达到最大重试次数，返回超时提示")
                return jsonify({
                    "content": "抱歉，响应时间较长，请稍后再试。",
                    "role": "assistant",
                    "session_id": session_id,
                    "conversation_id": conversation_id,
                    "warning": "获取完整回复超时"
                })
            
            # 如果没有获取到chat_id，尝试从消息事件中获取回复
            print("尝试从消息事件中获取回复")
            full_content = ""
            for event in events:
                if not event.strip():
                    continue
                
                event_lines = event.strip().split("\n")
                event_type = None
                event_data = None
                
                for line in event_lines:
                    if line.startswith("event:"):
                        event_type = line[6:].strip()
                    elif line.startswith("data:"):
                        event_data = line[5:].strip()
                
                # 如果是消息事件，尝试提取回复内容
                if event_type == "conversation.message.delta" and event_data:
                    try:
                        data_json = json.loads(event_data)
                        if data_json.get("role") == "assistant":
                            # 获取内容并累加
                            content_part = data_json.get("content", "")
                            if content_part:
                                full_content += content_part
                                print(f"从delta事件中获取部分内容: {content_part}")
                    except Exception as e:
                        print(f"解析消息事件失败: {str(e)}")
            
            # 如果累积了内容，直接返回
            if full_content:
                print(f"从消息事件中获取累积回复: {full_content[:100]}...")  # 只打印前100个字符
                return jsonify({
                    "content": full_content,
                    "role": "assistant",
                    "session_id": session_id,
                    "conversation_id": conversation_id
                })
            
            # 如果仍然没有获取到回复，尝试直接解析原始响应
            print("尝试直接解析原始响应")
            try:
                # 检查原始响应是否包含有用信息
                if "content" in chat_response.text:
                    import re
                    # 尝试使用正则表达式提取内容
                    content_match = re.search(r'"content"\s*:\s*"([^"]*)"', chat_response.text)
                    if content_match:
                        extracted_content = content_match.group(1)
                        print(f"从原始响应中提取内容: {extracted_content[:100]}...")
                        return jsonify({
                            "content": extracted_content,
                            "role": "assistant",
                            "session_id": session_id,
                            "conversation_id": conversation_id
                        })
            except Exception as e:
                print(f"解析原始响应失败: {str(e)}")
            
            # 如果仍然没有获取到回复，返回友好提示
            print("未能获取聊天回复，返回友好提示")
            return jsonify({
                "content": "抱歉，我暂时无法处理您的请求，请稍后再试。",
                "role": "assistant",
                "session_id": session_id,
                "conversation_id": conversation_id,
                "warning": "未能获取聊天回复"
            })
            
        except Exception as e:
            import traceback
            error_traceback = traceback.format_exc()
            print(f"处理响应失败: {str(e)}")
            print(f"错误详情: {error_traceback}")
            
            # 返回友好的错误信息
            return jsonify({
                "content": "抱歉，处理您的请求时出现了问题，请稍后再试。",
                "role": "assistant",
                "session_id": session_id,
                "conversation_id": conversation_id,
                "error": str(e)
            }), 500
            
    except Exception as e:
        import traceback
        error_traceback = traceback.format_exc()
        print(f"处理请求失败: {str(e)}")
        print(f"错误详情: {error_traceback}")
        
        return jsonify({
            "content": "抱歉，服务器处理您的请求时出现了问题，请稍后再试。",
            "role": "assistant",
            "error": str(e)
        }), 500


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8093, debug=True)
